import os
import pickle
import sys

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt
import numpy as np

results_directory = sys.argv[1]
test_name_suffix = sys.argv[2]
output_file_name = os.path.join(results_directory, sys.argv[3])
loads = [int(x) for x in sys.argv[4].split(",")]
apps = [x for x in sys.argv[5].split(",")]
assert len(apps) == 1

# usage：accuracy_per_process[(method, process_id)] = acc
with open(
        results_directory + "/process_acc_" + apps[0] + "_" + test_name_suffix + "_" + str(loads[0]) + "_1_1_0.0.pickle",
        'rb') as afile:
    accuracy_per_process = pickle.load(afile)
    _accuracy_per_process = {}
    for (method, process_id), acc in accuracy_per_process.items():
        if method == "MaxScoreBatchSubsetWithSkips":
            method = "TraceWeaver"
        _accuracy_per_process[(method, process_id)] = acc
    accuracy_per_process = _accuracy_per_process

methods = sorted(set(method for method, _ in accuracy_per_process.keys()))
processes = sorted(set(pid for _, pid in accuracy_per_process.keys()))

accuracy_matrix = np.zeros((len(methods), len(processes)))

for (method, process_id), acc in accuracy_per_process.items():
    row = methods.index(method)
    col = processes.index(process_id)
    accuracy_matrix[row, col] = acc

fig, ax = plt.subplots()
cax = ax.matshow(accuracy_matrix, cmap='Blues')
fig.colorbar(cax)

ax.set_xticklabels([''] + processes)
ax.set_yticklabels([''] + methods)

ax.xaxis.set_major_locator(plt.MultipleLocator(1))
ax.yaxis.set_major_locator(plt.MultipleLocator(1))

for (i, j), acc in np.ndenumerate(accuracy_matrix):
    ax.text(j, i, f'{acc:.2f}', ha='center', va='center')

plt.title(f'Accuracy per Process for {apps[0]}')
plt.savefig(output_file_name)
